`include "defines.v"

module WB_Stage (
    input  wire [ 4: 0] rd_addr_i,
    output wire [ 4: 0] rd_addr_o,

    input  wire         is_system_inst,
    input  wire [11: 0] csr_addr_i,
    output wire         csr_wen,
    output wire [11: 0] csr_addr_o,

    input  wire [`xlen] data_from_mem,
    input  wire [`xlen] ALU_out,
    output wire [`xlen] data_to_reg,
    output wire [`xlen] csr_data,

    input  wire         MemToReg,
    input  wire         RegWr_i,
    input  wire [ 2: 0] load_ext_type,
    output wire         RegWr_o
);

reg  [`xlen] load_ext_data;

always @(*) begin
    case (load_ext_type)
        `SEXT8 :  begin
            load_ext_data = { {56{data_from_mem[7]}} , data_from_mem[7:0] };
        end
        `SEXT16 : begin
            load_ext_data = { {48{data_from_mem[15]}} , data_from_mem[15:0] };
        end
        `SEXT32 : begin
            load_ext_data = { {32{data_from_mem[31]}} , data_from_mem[31:0] };
        end
        `ZEXT8 :  begin
            load_ext_data = { 56'b0 , data_from_mem[7:0] };
        end
        `ZEXT16 : begin
            load_ext_data = { 48'b0 , data_from_mem[15:0] };
        end
        `ZEXT32 : begin
            load_ext_data = { 32'b0 , data_from_mem[31:0] };
        end
        default : load_ext_data = data_from_mem;
    endcase
end

assign data_to_reg = MemToReg ? load_ext_data : ALU_out;

assign csr_addr_o = csr_addr_i;
assign csr_data = data_from_mem;

assign rd_addr_o = rd_addr_i;
assign RegWr_o = RegWr_i;

endmodule